# 同agent_cdm_2.py
# 仅仅修改了    if answer_content.strip().lower() == ground_truth.strip().lower():

import re
import random
from word2number import w2n

def convert_word_number(text: str) -> str:
    try:
        text = str(w2n.word_to_num(text))
    except Exception:
        pass
    return text


def _fix_fracs(string):
    substrs = string.split('\\frac')
    new_str = substrs[0]
    if len(substrs) > 1:
        substrs = substrs[1:]
        for substr in substrs:
            new_str += '\\frac'
            if len(substr) > 0 and substr[0] == '{':
                new_str += substr
            else:
                try:
                    assert len(substr) >= 2
                except Exception:
                    return string
                a = substr[0]
                b = substr[1]
                if b != '{':
                    if len(substr) > 2:
                        post_substr = substr[2:]
                        new_str += '{' + a + '}{' + b + '}' + post_substr
                    else:
                        new_str += '{' + a + '}{' + b + '}'
                else:
                    if len(substr) > 2:
                        post_substr = substr[2:]
                        new_str += '{' + a + '}' + b + post_substr
                    else:
                        new_str += '{' + a + '}' + b
    string = new_str
    return string


def _fix_a_slash_b(string):
    if len(string.split('/')) != 2:
        return string
    a = string.split('/')[0]
    b = string.split('/')[1]
    try:
        if 'sqrt' not in a:
            a = int(a)
        if 'sqrt' not in b:
            b = int(b)
        assert string == '{}/{}'.format(a, b)
        new_string = '\\frac{' + str(a) + '}{' + str(b) + '}'
        return new_string
    except Exception:
        return string


def _fix_sqrt(string):
    _string = re.sub(r'\\sqrt(\w+)', r'\\sqrt{\1}', string)
    return _string

def strip_answer_string(string):
    string = str(string).strip()
    # linebreaks
    string = string.replace('\n', '')

    # right "."
    string = string.rstrip('.')

    # remove inverse spaces
    # replace \\ with \
    string = string.replace('\\!', '')
    # string = string.replace("\\ ", "")
    # string = string.replace("\\\\", "\\")

    # matrix
    string = re.sub(r'\\begin\{array\}\{.*?\}', r'\\begin{pmatrix}', string)
    string = re.sub(r'\\end\{array\}', r'\\end{pmatrix}', string)
    string = string.replace('bmatrix', 'pmatrix')

    # replace tfrac and dfrac with frac
    string = string.replace('tfrac', 'frac')
    string = string.replace('dfrac', 'frac')
    string = (string.replace('\\neq', '\\ne').replace('\\leq', '\\le').replace('\\geq', '\\ge'))

    # remove \left and \right
    string = string.replace('\\left', '')
    string = string.replace('\\right', '')
    string = string.replace('\\{', '{')
    string = string.replace('\\}', '}')

    # Function to replace number words with corresponding digits
    def replace_match(match):
        word = match.group(1).lower()
        if convert_word_number(word) == word:
            return match.group(0)
        else:
            return convert_word_number(word)

    string = re.sub(r'\\text\{([a-zA-Z]+)\}', replace_match, string)

    # Before removing unit, check if the unit is squared (for surface area)
    string = re.sub(r'(cm|inches)\}\^2', r'\1}', string)

    # Remove unit: miles, dollars if after is not none
    _string = re.sub(r'\\text{.*?}$', '', string).strip()
    if _string != '' and _string != string:
        # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string))
        string = _string

    # Remove circ (degrees)
    string = string.replace('^{\\circ}', '')
    string = string.replace('^\\circ', '')

    # remove dollar signs
    string = string.replace('\\$', '')
    string = string.replace('$', '')
    string = string.replace('\\(', '').replace('\\)', '')

    # convert word number to digit
    string = convert_word_number(string)

    # replace "\\text{...}" to "..."
    string = re.sub(r'\\text\{(.*?)\}', r'\1', string)
    for key in ['x=', 'y=', 'z=', 'x\\in', 'y\\in', 'z\\in', 'x\\to', 'y\\to', 'z\\to']:
        string = string.replace(key, '')
    string = string.replace('\\emptyset', r'{}')
    string = string.replace('(-\\infty,\\infty)', '\\mathbb{R}')

    # remove percentage
    string = string.replace('\\%', '')
    string = string.replace('\%', '')
    string = string.replace('%', '')

    # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string
    string = string.replace(' .', ' 0.')
    string = string.replace('{.', '{0.')

    # cdot
    # string = string.replace("\\cdot", "")
    if (string.startswith('{') and string.endswith('}') and string.isalnum()
            or string.startswith('(') and string.endswith(')') and string.isalnum()
            or string.startswith('[') and string.endswith(']') and string.isalnum()):
        string = string[1:-1]

    # inf
    string = string.replace('infinity', '\\infty')
    if '\\infty' not in string:
        string = string.replace('inf', '\\infty')
    string = string.replace('+\\inity', '\\infty')

    # and
    string = string.replace('and', '')
    string = string.replace('\\mathbf', '')

    # use regex to remove \mbox{...}
    string = re.sub(r'\\mbox{.*?}', '', string)

    # quote
    string.replace("'", '')
    string.replace('"', '')

    # i, j
    if 'j' in string and 'i' not in string:
        string = string.replace('j', 'i')

    # replace a.000b where b is not number or b is end, with ab, use regex
    string = re.sub(r'(\d+)\.0*([^\d])', r'\1\2', string)
    string = re.sub(r'(\d+)\.0*$', r'\1', string)

    # if empty, return empty string
    if len(string) == 0:
        return string
    if string[0] == '.':
        string = '0' + string

    # to consider: get rid of e.g. "k = " or "q = " at beginning
    if len(string.split('=')) == 2:
        if len(string.split('=')[0]) <= 2:
            string = string.split('=')[1]

    string = _fix_sqrt(string)
    string = string.replace(' ', '')

    # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b}
    string = _fix_fracs(string)

    # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y
    string = _fix_a_slash_b(string)

    # Remove unnecessary '\' before integers
    string = re.sub(r'\\(?=\-?\d+(\\|\)|,|\]|$))', '', string)

    # Remove grade level (e.g., 12th grade) and just maintain the integer
    string = re.sub(r'thgrade$', '', string)

    # If the answer is a list of integers (without parenthesis), sort them
    if re.fullmatch(r'(\s*-?\d+\s*,)*\s*-?\d+\s*', string):
        # Split the string into a list of integers
        try:
            integer_list = list(map(int, string.split(',')))
        except Exception:
            integer_list = list(map(int, '-1,-1'.split(',')))

        # Sort the list in ascending order
        sorted_list = sorted(integer_list)

        # Join the sorted list back into a comma-separated string
        string = ','.join(map(str, sorted_list))

    return string

def extract_answer(pred_str, use_last_number=True):
    pred_str = pred_str.replace('\u043a\u0438', '')
    if 'final answer is $' in pred_str and '$. I hope' in pred_str:
        # minerva_math
        tmp = pred_str.split('final answer is $', 1)[1]
        pred = tmp.split('$. I hope', 1)[0].strip()
    elif 'boxed' in pred_str:
        ans = pred_str.split('boxed')[-1]
        if len(ans) == 0:
            return ''
        elif ans[0] == '{':
            stack = 1
            a = ''
            for c in ans[1:]:
                if c == '{':
                    stack += 1
                    a += c
                elif c == '}':
                    stack -= 1
                    if stack == 0:
                        break
                    a += c
                else:
                    a += c
        else:
            a = ans.split('$')[0].strip()
        pred = a
    elif 'he answer is' in pred_str:
        pred = pred_str.split('he answer is')[-1].strip()
    elif 'final answer is' in pred_str:
        pred = pred_str.split('final answer is')[-1].strip()
    elif '答案是' in pred_str:
        # Handle Chinese few-shot multiple choice problem answer extraction
        pred = pred_str.split('答案是')[1].strip().split('\n\n')[0].strip()
    else:  # use the last number
        if use_last_number:
            pattern = '-?\d*\.?\d+'
            pred = re.findall(pattern, pred_str.replace(',', ''))
            if len(pred) >= 1:
                pred = pred[-1]
            else:
                pred = ''
        else:
            pred = ''

    # multiple line
    # pred = pred.split("\n")[0]
    pred = re.sub(r'\n\s*', '', pred)
    if pred != '' and pred[0] == ':':
        pred = pred[1:]
    if pred != '' and pred[-1] == '.':
        pred = pred[:-1]
    if pred != '' and pred[-1] == '/':
        pred = pred[:-1]
    pred = strip_answer_string(pred)
    return pred

choices = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']

class ResponseParser:

    @staticmethod
    def parse_first_capital(text: str, options: list[str]) -> str:
        for t in text:
            if t.isupper() and (t in options):
                return t
        return ''

    @staticmethod
    def parse_last_capital(text: str, options: list[str]) -> str:
        for t in text[::-1]:
            if t.isupper() and (t in options):
                return t
        return ''

    @staticmethod
    def parse_first_option_with_choices(text: str, options: list[str]) -> str:
        """
        Find first valid option for text.

        Args:
            text: The text to parse.
            options: The options to find. e.g. ['A', 'B', 'C', 'D']
        """
        options_concat = ResponseParser.process_options(options)

        patterns = [
            rf'答案是?\s?([{options_concat}])',
            rf'答案是?\s?：([{options_concat}])',
            rf'答案是?\s?:([{options_concat}])',
            rf'答案应该?是\s?([{options_concat}])',
            rf'答案应该?选\s?([{options_concat}])',
            rf'答案为\s?([{options_concat}])',
            rf'答案选\s?([{options_concat}])',
            rf'选择?\s?([{options_concat}])',
            rf'故选?\s?([{options_concat}])'
            rf'只有选?项?\s?([{options_concat}])\s?是?对',
            rf'只有选?项?\s?([{options_concat}])\s?是?错',
            rf'只有选?项?\s?([{options_concat}])\s?不?正确',
            rf'只有选?项?\s?([{options_concat}])\s?错误',
            rf'说法不?对选?项?的?是\s?([{options_concat}])',
            rf'说法不?正确选?项?的?是\s?([{options_concat}])',
            rf'说法错误选?项?的?是\s?([{options_concat}])',
            rf'([{options_concat}])\s?是正确的',
            rf'([{options_concat}])\s?是正确答案',
            rf'选项\s?([{options_concat}])\s?正确',
            rf'所以答\s?([{options_concat}])',
            rf'所以\s?([{options_concat}][.。$]?$)',
            rf'所有\s?([{options_concat}][.。$]?$)',
            rf'[\s，：:,]([{options_concat}])[。，,\.]?$',
            rf'[\s，,：:][故即]([{options_concat}])[。\.]?$',
            rf'[\s，,：:]因此([{options_concat}])[。\.]?$',
            rf'[是为。]\s?([{options_concat}])[。\.]?$',
            rf'因此\s?([{options_concat}])[。\.]?$',
            rf'显然\s?([{options_concat}])[。\.]?$',
            rf'答案是\s?(\S+)(?:。|$)',
            rf'答案应该是\s?(\S+)(?:。|$)',
            rf'答案为\s?(\S+)(?:。|$)',
            rf'答案是(.*?)[{options_concat}]',
            rf'答案为(.*?)[{options_concat}]',
            rf'固选(.*?)[{options_concat}]',
            rf'答案应该是(.*?)[{options_concat}]',
            rf'[Tt]he answer is \(?[{options_concat}]\)?',
            rf'[Tt]he correct answer is [{options_concat}]',
            rf'[Tt]he correct answer is:\n[{options_concat}]',
            rf'(\s|^)[{options_concat}][\s。，,\.$]',  # noqa
            rf'^选项\s?([{options_concat}])',
            rf'^([{options_concat}])\s?选?项',
            rf'(\s|^)[{options_concat}][\s。，,：:\.$]',
            rf'(\s|^)[{options_concat}](\s|$)',
            rf'[{options_concat}]',
        ]

        regexes = [re.compile(pattern) for pattern in patterns]
        for regex in regexes:
            match = regex.search(text)
            if match:
                outputs = match.group(0)
                for i in options:
                    if i in outputs:
                        return i
        # If no match found, try to find the last capital letter in the text
        last_capital = ResponseParser.parse_last_capital(text, options)
        if last_capital:
            return last_capital
        return 'No valid option found'

    @staticmethod
    def parse_first_option(text: str, options: list[str]) -> str:
        """
        Find first valid option for text.

        Args:
            text: The text to parse.
        """
        options_pattern = ResponseParser.process_options(options)

        patterns = [
            rf'[Aa]nswer:\s*({options_pattern})',
            rf'ANSWER:\s*({options_pattern})',
            rf'answer is \(?({options_pattern})\)?',
            rf'[Tt]he correct answer is:\s*({options_pattern})',
            rf'[Tt]he correct answer is:\n\s*({options_pattern})',
            rf'[Tt]he correct answer is:\n\n-\s*({options_pattern})',
            rf'[Tt]he answer might be:\n\n-\s*({options_pattern})',
            rf'[Tt]he answer is \s*({options_pattern})',
        ]

        regexes = [re.compile(pattern) for pattern in patterns]
        for regex in regexes:
            matches = regex.search(text)
            if matches:
                return matches.group(1)
        # If no match found, try to find the last capital letter in the text
        last_capital = ResponseParser.parse_last_capital(text, options)
        if last_capital:
            return last_capital
        return 'No valid option found'


    @staticmethod
    def parse_bracketed_answer(text: str, options: list[str]) -> str:
        options = ResponseParser.process_options(options)
        # Match the first occurrence of the options in angle brackets
        match = re.search(rf'<({options})>', text)
        if match:
            return match.group(1)
        return 'No valid option found'

    @staticmethod
    def process_options(options: list[str]) -> str:
        # Escape each option to ensure special characters in options are treated literally
        escaped_options = [re.escape(option) for option in options]
        # Join options into a regex pattern separated by '|', to match any of the options
        options_pattern = '|'.join(escaped_options)
        return options_pattern

def extract_box_answer(answer_str):
    if 'boxed' in answer_str:
        ans = answer_str.split('boxed')[-1]
        if len(ans) == 0:
            return ''
        elif ans[0] == '{':
            stack = 1
            a = ''
            for c in ans[1:]:
                if c == '{':
                    stack += 1
                    a += c
                elif c == '}':
                    stack -= 1
                    if stack == 0:
                        break
                    a += c
                else:
                    a += c
            return a
        else:
            return answer_str
    else:
        return answer_str


def extract_answer_content(solution_str,ground_truth):
    # Remove everything before the first "Assistant:"
    # if "Assistant:" in solution_str:
    #     solution_str = solution_str.split("Assistant:", 1)[1]
    # else:
    #     return None

    
    # solution_str= solution_str.split('</think>')[-1].strip()
    # final_answer = extract_answer(solution_str)
    # if final_answer == '':
    #         final_answer = None
    # return final_answer
    
    answer_pattern = r'<answer>(.*?)</answer>'
    matches = re.findall(answer_pattern, solution_str, re.DOTALL)

    if matches:
        answer_content = matches[-1].strip()
        if answer_content == '':
            answer_content = None
    else:
        answer_content = None
    
    
        
    if answer_content is not None:
        answer_content = extract_box_answer(answer_content)
        if ground_truth.upper() in choices:
            answer_content = ResponseParser.parse_first_option(answer_content, options=choices)
        else:
            answer_content = extract_answer(answer_content)
            if answer_content == '':
                answer_content = None

    # answer_content = ResponseParser.parse_first_option(solution_str, options=choices)
    return answer_content

def check_agent_str(solution_str):
    if "Agent 1" in solution_str or "Agent 2" in solution_str or "Agent 3" in solution_str:
        return True
    else:
        return False
    
def check_format(solution_str):
    answer_count = solution_str.count("<answer>")
    think_count = solution_str.count("</think>")
    solution_list = solution_str.split("\n")
    
    if answer_count == 1 and think_count == 1:
        answer_str = solution_str.split("</think>")[-1].strip()
        if answer_str.startswith("<answer>") and "</answer>" in solution_list[-1]:
            return True
        else:
            return False
    else:
        return False

def score_sum(score_dict):
    """Sum the scores in the score_dict."""
    total_score = 0
    for key, value in score_dict.items():
        if key == 'score':
            continue
        if isinstance(value, (int, float)):
            total_score += value
        elif isinstance(value, dict):
            total_score += score_sum(value)
    return total_score
    

def compute_score(data_source, solution_str, ground_truth, extra_info, method='strict', format_score=0.1, score=1.):
    """The scoring function for GSM8k.

    Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

    Args:
        solution_str: the solution text
        ground_truth: the ground truth
        method: the method to extract the solution, choices are 'strict' and 'flexible'
        format_score: the score for the format
        score: the score for the correct answer
    """
    score_dict = {
        "score": 0,
        "format_score": 0,
        "agent_str_score": 0,
        "answer_score": 0,
    }
    answer_content = extract_answer_content(solution_str=solution_str,ground_truth=ground_truth)
    do_print = random.randint(1, 64) == 1
    if do_print:
        print(f"--------------------------------")
        print(f"Ground truth: {ground_truth} | Extracted answer: {answer_content}")
        print(f"Solution string: {solution_str}")

    if answer_content is None:
        if do_print:
            print(f"No answer found")
        return score_dict
    
    score_dict["format_score"] = format_score

    if check_format(solution_str):
        score_dict["format_score"] += format_score


    if answer_content.strip().lower() == ground_truth.strip().lower():
        if do_print:
            print(f"Correct answer: {answer_content}")

        score_dict["answer_score"] = score
        score_dict["score"] = score_sum(score_dict)

        return score_dict
    else:
        if do_print:
            print(f"Incorrect answer {answer_content} | Ground truth: {ground_truth}")
            
        score_dict["score"] = score_sum(score_dict)

        return score_dict
